import os
import json
import datetime
import argparse

import torch
import wandb

parser = argparse.ArgumentParser(description="训练TransE图嵌入")
parser.add_argument("-g", "--graph", type=str, default="data/Graph", help="要训练的图")
parser.add_argument("-b", "--batch_size", type=int, default=1024, help="每批训练样本的数量")
parser.add_argument("-e", "--epochs", type=int, default=10000, help="训练的总回合数")
parser.add_argument("-lr", "--learning_rate", type=float, default=0.001, help="学习率")
parser.add_argument("-vf", "--validate_frequency", type=float, default=100, help="每多少个Epoch验证一次性能")
parser.add_argument("-n", "--num_checkpoints", type=int, default=2, help="保存模型的最大次数")
parser.add_argument("-o", "--outputs", type=str, default="outputs/TransE", help="输出权重路径")
parser.add_argument("-r", "--resume", type=bool, default=True, help="输出权重路径")

if __name__ == "__main__":
    args = parser.parse_args()
    timestamp = datetime.datetime.strftime(datetime.datetime.now(), r"%Y%m%d-%H%M%S")
    logger = wandb.init(project="TransE", name=timestamp)

    from dataset.graph import Graph
    g = Graph(args.graph)

    from dataset.dataset import GraphDataset
    mix_dataset = GraphDataset(g)

    from model.transe import TransE
    model = TransE(g.num_entities, g.num_relations, 128)

    weights = None
    if os.path.isfile(os.path.join(args.outputs, "latest_checkpoints.json")):
        with open(os.path.join(args.outputs, "latest_checkpoints.json"), 'r', encoding="utf-8") as fp:
            last_ckpt_names = json.load(fp)
        if last_ckpt_names:
            if os.path.isfile(last_ckpt_names[0][1]):
                weights = torch.load(last_ckpt_names[0][1])


    from trainer.trainer import Trainer
    trainer = Trainer(model, mix_dataset, mix_dataset, margin=1, \
                      lr=args.learning_rate, batch_size=args.batch_size, \
                      weights=weights, max_ckpt=2, outputs=args.outputs)
    
    for e in range(args.epochs):
        loss = trainer.train(e)
        # logger.log({"loss": loss})
        if (e+1) % args.validate_frequency == 0:
            metrics = trainer.validate(e)
            logger.log({"loss": loss, "MRR": metrics["MRR"], "hit@5":metrics["hit@5"]})


